from __future__ import print_function

import json
import copy


class TrainingMetadata(object):

    def __init__(self):
        self.batch_size = 16
        self.nb_epoch = 1
        self.nb_iter = 1
        self.files_per_batch = 1
        self.optimizer = None
        self.shuffle = True
        self.performance = {'systole': dict(), 'diastole': dict(), 'crps': [], 'val_crps': []}
        self.loss = ''

    def to_json(self):
        training_metadata = copy.copy(self)
        training_metadata.optimizer = self.optimizer.get_config()
        meta_json = json.dumps(training_metadata, default=lambda o: o.__dict__, indent=4)
        return meta_json

    def save_to_json(self, path):
        # save metadata to json
        f = open(path, 'wb')
        training_json = self.to_json()
        f.write(training_json)
        f.close()

    def update_performance(self, hist_systole, hist_diastole, crps, val_crps):
        for k, v in hist_systole.items():
            if k == 'batch' or k == 'size':
                continue
            if k not in self.performance['systole']:
                self.performance['systole'][k] = []
            if k not in self.performance['diastole']:
                self.performance['diastole'][k] = []

            self.performance['systole'][k] += hist_systole[k]
            self.performance['diastole'][k] += hist_diastole[k]
        self.performance['crps'] += crps
        self.performance['val_crps'] += val_crps
        return self.performance

